Summary:

• Refactored io_manager into five distinct components:
  • DecoderRunner: Module wrapper class.
  • PromptProcessor: Handles prompt processing using the decoder and key-value manager.
  • TokenGenerator: Generates tokens using the decoder and key-value manager.
  • KVManager: Manages key-value cache with kv_updater, including data buffer allocation, cache updates, and buffer updates in TensorImpl.
  • IBufferAlloc: Allocates data buffers from RPC memory or client buffer.
  • Validated story llama with CL=128, prefill_ar_len=16, QNN SDK: 2.32
  • Original :

• Refactor:

• Support multi-turn use case.
  • Validated on story llama. To simulate the scenario, I forced decode mode to generate 5 tokens each time. Tokens with random length are inserted after one round of prefill->decode finished.
  • Reproduce command: (Note that some whitespaces are missing due to decoding. But token is actually the same as golden.)

  python examples/qualcomm/oss_scripts/llama/llama.py -b build-android --checkpoint stories110M.pt --params params.json --tokenizer_model tokenizer.model --prompt "Once" "a little girl named Lily." "toys and her favorite toy was a big, red ball." "s mom asked her to help her with the laundry." "and she put all the clothes in the washing machine." --temperature 0 --tokenizer_bin tokenizer.bin --llama_model stories110m --model_mode hybrid --ptq 16a4w -m SM8650 -H ${HOST} -s ${DEVICE}-a ${ARTIFACTS}--max_seq_len 128 --prefill_ar_len 16
  Result:
  Once upon a time, there wasa little girl named Lily. She loved to play with hertoys and her favorite toy was a big, red ball. One day, Lily's mom asked her to help her with the laundry. Lily was happy to helpand she put all the clothes in the washing machine.
  After the clothes were
  

  • Need to apply the below patch to forced decode mode to generate 5 tokens each time.

diff --git a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
index 8a81b598d..a8ec53cdb 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
+++ b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
@@ -170,7 +170,10 @@ Result<int64_t> TokenGenerator::generate(
       "Failed to set output tensor for module %s",
       forward_name_.c_str());
   // Generate our tokens
-  while (pos < seq_len - 1) {
+  // force decode to generate 5 runs at most
+  int64_t max_pos = std::min(pos + 5, (int64_t)seq_len - 1);
+//   while (pos < seq_len - 1) {
+  while (pos < max_pos) {
     // Fill in the token and position data
     prepare_io(cur_token, pos);
     // Only update data pointer of the cache to the tensor for SHIFT_POINTER